from typing import List

from src.common.config import config
from src.common.logger import getLogger
from langchain_core.documents import Document
from langchain_qdrant import QdrantVectorStore, SparseEmbeddings, SparseVector
from qdrant_client import QdrantClient, models
from langchain_text_splitters import RecursiveCharacterTextSplitter

logger = getLogger()

class HybridStore:

    def __init__(self):
        config_dict = config.parse_config_key(["qdrant"])
        logger.info(f"HybridStore __init__ config_dict: {config_dict}")
        self.host = config_dict.get("host", "")
        self.port = config_dict.get("port", 0)
        self.timeout = config_dict.get("timeout", 0)
        self.vector_size = config_dict.get("vector_size", 0)

    def new_vector_client(self):
        qdrant_client = QdrantClient(
            host = self.host,
            port = self.port,
            timeout = self.timeout
        )
        logger.info(f"HybridStore new_vector_client qdrant_client: {qdrant_client}")
        return qdrant_client

    def new_vector_store(self, embedding, collection_name):
        logger.info(f"HybridStore new_vector_store collection_name: {collection_name}")
        qdrant_client = self.new_vector_client()
        exist = qdrant_client.collection_exists(collection_name)
        logger.info(f"HybridStore new_vector_store collection_name exist: {exist}")
        if not exist:
            qdrant_client.create_collection(
                collection_name = collection_name,
                vectors_config = models.VectorParams(size = self.vector_size, distance = models.Distance.COSINE),
                sparse_vectors_config = {
                    "spare_vectors": models.SparseVectorParams(
                        index = models.SparseIndexParams(full_scan_threshold = 1000, on_disk = True, datatype = "uint8"),
                        modifier = None
                    )
                }
            )
        logger.info(f"HybridStore new_vector_store collection_name info: {qdrant_client.info()}")

        vector_store = QdrantVectorStore(
            client = qdrant_client,
            collection_name = collection_name,
            embedding = embedding,
            retrieval_mode = "dense",
            sparse_embedding = embedding,
            sparse_vector_name = "spare_vectors"
        )
        return vector_store

    def add_document_vector(self, embedding, collection_name, content):
        logger.info(f"HybridStore add_document_vector collection_name: {collection_name}")
        vector_store = self.new_vector_store(embedding, collection_name)
        logger.info(f"HybridStore add_document_vector vector_store {collection_name} finished")

        text_splitter = RecursiveCharacterTextSplitter(chunk_size = 768, chunk_overlap = 50)
        split_texts = text_splitter.split_text(content)
        documents = [Document(page_content = text) for text in split_texts]
        logger.info(f"HybridStore add_document_vector documents len: {len(documents)}")

        vector_store.add_documents(documents)
        logger.info(f"HybridStore add_document_vector vector_store {collection_name} finished")

class CustomSparseEmbed(SparseEmbeddings):

    def __init__(self, embed_model):
        super().__init__()
        self.embed_model = embed_model

    def embed_documents(self, texts: List[str]) -> List[SparseVector]:
        if not texts or texts is None:
            return None

        sparse_vecs = []
        for text in texts:
            sparse_vecs.append(self.embed_query(text))
        return sparse_vecs

    def embed_query(self, text: str) -> SparseVector:
        if not text or text is None:
            return None

        output = self.embed_model.encode(
            sentences = text,
            batch_size = 12,
            max_length = 8192,
            return_dense = True,
            return_sparse = True,
            return_colbert_vecs = False
        )
        sparse_vectors = output['lexical_weights']

        indices, values = [], []
        for key, value in sparse_vectors.items():
            indices.append(key)
            values.append(value)
        return SparseVector(indices = indices, values = values)
